package com.alisonyu.spider.Downloader;

import com.alisonyu.spider.Spider;
import org.openqa.selenium.TimeoutException;
import org.openqa.selenium.WebDriver;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * @author yuzhiyi
 * @date 2018/7/19 14:51
 */
public class BrowserDownloader implements Download{

    private WebDriver webDriver;
    private int maxRetryTime;

    public BrowserDownloader(){
        //设置默认WebDriver为Chrome
        this.webDriver = BrowserUtil.getChrome();
        //设置超时时间为10秒
        this.webDriver.manage().timeouts().pageLoadTimeout(10, TimeUnit.SECONDS);
    }

    @Override
    public CompletableFuture<HttpResponse> asyncDownload(Request request) {
        CompletableFuture<HttpResponse> future = new CompletableFuture<>();
        Spider.handleTask(()->{
            future.complete(download(request));
        });
        return future;
    }

    @Override
    public HttpResponse download(Request request) {
        String url = request.getUrl();
        try{
            webDriver.get(url);
            String content = webDriver.getPageSource();
            HttpResponse response = new HttpResponse();
            response.setContent(content);
            response.setStatusCode(200);
            response.setSuccess(true);
            response.setBytes(content.getBytes());
            return response;
        }catch (TimeoutException e){
            if(request.getRetryTime()<maxRetryTime){
                request.setRetryTime(request.getRetryTime()+1);
                return download(request);
            }
        }catch (Exception e){
            e.printStackTrace();
        }
        HttpResponse resp = new HttpResponse();
        resp.setSuccess(false);
        return resp;
    }

    @Override
    public void setMaxRetryTime(int k) {
        this.maxRetryTime = k;
    }
}
